package com.utils;

import com.github.unidbg.AndroidEmulator;
import com.github.unidbg.Emulator;
import com.github.unidbg.Module;
import com.github.unidbg.ModuleListener;
import com.github.unidbg.arm.backend.*;

import java.util.*;

public class SearchRc4 implements ModuleListener{
    public Module targetModule;
    public AndroidEmulator emulator;
    public String moduleName;
    public int intervals;
    public long blockcount = 0;
    public long startcount = 0;
    public List<long[]> activePlace = new ArrayList<>();;
    public SortedSet<Long> heapAddress = new TreeSet();
    public ArrayList<byte[]> memoryRange = new ArrayList<>();


    @Override
    public void onLoaded(Emulator<?> emulator, Module module) {
        if(module.name.equals(moduleName)){
            this.emulator = (AndroidEmulator) emulator;
            targetModule = module;
            hookReadAndWrite();
            hookBlock();
        }
    }


    public SearchRc4(String moduleName, int intervals){
        this.moduleName = moduleName;
        this.intervals = intervals;
        this.startcount = 0;
    };


    public void setMemoryRange(int minSize) {
        memoryRange = new ArrayList<>();;
        // 整个栈空间
        long stackend = 0xC0000000L;
        long stackSize = 256L * emulator.getPageAlign();
        memoryRange.add(emulator.getBackend().mem_read(stackend - stackSize, stackSize));

        // 搜索除栈外活跃的内存块
        summaryRanges();
        for(long[] one: activePlace){
            try
            {
                long start = one[0];
                long size = one[1] - one[0] + 1;
                if(size > minSize){
                    byte[] memory = emulator.getBackend().mem_read(start, size);
                    memoryRange.add(memory);
                }

            }catch(Exception ignored)
            {

            }

        }
    }

    public void hookReadAndWrite(){
        emulator.getBackend().hook_add_new(new WriteHook() {
            @Override
            public void hook(Backend backend, long address, int size, long value, Object user) {
                for(int i=0;i<size;i++){
                    heapAddress.add(address+i);
                }
            }

            @Override
            public void onAttach(UnHook unHook) {

            }

            @Override
            public void detach() {

            }
        }, emulator.getMemory().MMAP_BASE, 0x7fffffff, null);


        emulator.getBackend().hook_add_new(new ReadHook() {
            @Override
            public void hook(Backend backend, long address, int size, Object user) {
                for(int i=0;i<size;i++){
                    heapAddress.add(address+i);
                }
            }

            @Override
            public void onAttach(UnHook unHook) {

            }

            @Override
            public void detach() {

            }
        }, emulator.getMemory().MMAP_BASE, 0x7fffffff, null);
    }

    public void summaryRanges(){
        ArrayList<Long> nums = new ArrayList<Long>(heapAddress);
        activePlace = new ArrayList<>();
        long start=nums.get(0);
        long end;
        long temp;
        int i=0;


        while(i<nums.size()){
            temp=nums.get(i);
            // 如果下一个触底
            if(i+1 == nums.size()){
                end = nums.get(i);
                activePlace.add(new long[]{start, end});
                break;
            }

            // 如果不连续
            if((temp+1) != nums.get(i+1)){
                end = temp;
                activePlace.add(new long[]{start, end});
                start = nums.get(i+1);
            }
            i++;
        }
    }

    public void hookBlock(){
        emulator.getBackend().hook_add_new(new BlockHook() {
            @Override
            public void hookBlock(Backend backend, long address, int size, Object user) {
                blockcount += 1;
                if(startcount==0){
                    if(blockcount%intervals == 0){
                        setMemoryRange(0x100);
                        searchRc4InMemory();
                    }
                }

            }

            @Override
            public void onAttach(UnHook unHook) {

            }

            @Override
            public void detach() {

            }
        }, targetModule.base, targetModule.size+ targetModule.base, null);

    }


    public void searchRc4InMemory(){
        System.out.println("search Rc4");
        int rc4Length = 0x100;
        for(int i =0;i<memoryRange.size();i++){
            byte[] block = memoryRange.get(i);
            for(int s=0;s<block.length - rc4Length;s++){
                if(verifyRc4(Arrays.copyOfRange(block, s , s + rc4Length))){
                    System.out.println("find RC4!");
                    emulator.attach().debug();
                    break;
                };
            }
        }

    }

    public static Boolean verifyRc4(byte[] bytes) {
        HashSet<Byte> set = new HashSet<>();
        for (byte aByte : bytes) {
            set.add(aByte);
        }
        return set.size() == bytes.length;
    };



}
